package com.xcj.springai.config;


import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.NodeOutput;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.codec.ServerSentEvent;
import reactor.core.publisher.Sinks;

import java.util.Map;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Consumer;

/**
 * 作用： 异步执行图流程并将响应结果通过 SSE（Server-Sent Events）推送给前端。
 * 职责：
 * 遍历异步输出（AsyncGenerator<NodeOutput>）；
 * 将每个输出封装为 ServerSentEvent<String>；
 * 支持流式事件推送和异常处理。
 * 角色： 前后端之间的“桥梁”，流式处理图节点的输出。
 */

@Slf4j
public class GraphProcess {

    private static final Logger logger = LoggerFactory.getLogger(GraphProcess.class);

    private final ExecutorService executor = Executors.newSingleThreadExecutor();

    private CompiledGraph compiledGraph;

    public GraphProcess(CompiledGraph compiledGraph) {
        this.compiledGraph = compiledGraph;
    }
    public void processStream(AsyncGenerator<NodeOutput> generator, Sinks.Many<ServerSentEvent<String>> sink) {
        executor.submit(() -> {
            // 异步遍历generator中的每个输出
            generator.forEachAsync(output -> {
                try {
                    // 打印输出信息
                    logger.info("output = {}", output);
                    // 获取节点名称
                    String nodeName = output.node();
                    String content;
                    // 判断输出是否为StreamingOutput类型
                    if (output instanceof StreamingOutput streamingOutput) {
                        // 如果是StreamingOutput类型，将其转换为JSON字符串
                        content = JSON.toJSONString(Map.of(nodeName, streamingOutput.chunk()));
                    } else {
                        // 如果不是StreamingOutput类型，创建JSONObject并填充数据
                        JSONObject nodeOutput = new JSONObject();
                        nodeOutput.put("data", output.state().data());
                        nodeOutput.put("node", nodeName);
                        content = JSON.toJSONString(nodeOutput);
                    }
                    // 发送ServerSentEvent事件
                    sink.tryEmitNext(ServerSentEvent.builder(content).build());
                } catch (Exception e) {
                    // 捕获异常并抛出CompletionException
                    throw new CompletionException(e);
                }
            }).thenAccept(v -> {
                // 当异步操作正常完成时，发送完成事件
                sink.tryEmitComplete();
            }).exceptionally(e -> {
                // 捕获异步操作中的异常，并发送错误事件
                sink.tryEmitError(e);
                return null;
            });
        });
    }

}